module FES1DL21CTE_m
use mesh1DL2_m
use FES1DL2_m
use FES1DL21_m
use SMatrix_m

type, extends(FES1DL2) :: FES1DL21CTE
contains
    procedure :: initFES => initFES1DL21CTE
    
	procedure :: assembleHeatToElasticL2CTE_Func
	procedure :: assembleHeatToElasticL2CTE_Const
	procedure :: assembleHeatToElasticL2CTE_Vector
    generic :: assembleHeatToElastic => assembleHeatToElasticL2CTE_Func, assembleHeatToElasticL2CTE_Const, assembleHeatToElasticL2CTE_Vector
    
    procedure :: assembleHeatSourceL2CTE_const
    generic :: assembleHeatSource => assembleHeatSourceL2CTE_const
    
    procedure :: assembleDC1_L2CTE
    generic :: assembleDC1 => assembleDC1_L2CTE
    
    procedure :: assembleDC2_L2CTE_Const
    procedure :: assembleDC2_L2CTE_Function
    procedure :: assembleDC2_L2CTE_VM
    generic :: assembleDC2 => assembleDC2_L2CTE_Const,assembleDC2_L2CTE_Function, assembleDC2_L2CTE_VM
    
    procedure :: applyBCHeat => applyBCHeatL2CTE_function
    procedure :: applyBCZeroDisp => applyBCZeroDispL2CTE
    procedure :: extractUT => extractUTL2CTE
end type

contains
subroutine initFES1DL21CTE(this, mesh1D)
    implicit none
    class(FES1DL21CTE), intent(out) :: this
    type(Mesh1DL2), target, intent(in) :: mesh1D

    integer :: band
    this%dim = 1
    this%nDof = mesh1D%nv * 2
    this%nVertices = mesh1D%nv
    this%nElement = mesh1D%nt
    this%nBoundary = mesh1D%nb
    this%nPerElement = 2
    this%nPerBoundary = 1
    this%dofPerNode = 2
    this%dofPerElement = 4
    this%dofPerBoundary = 2
    
    band = mesh1D.elementDiff()
    this%nBand = (band + 1) * 2
    this%mesh1D => mesh1D
end subroutine

subroutine elementCoupledMatrixL2CTE(x, beta, tpCoupled, K)
    implicit none
    real*8, intent(in) :: x(2), beta(2)
    integer, intent(in) :: tpCoupled
    real*8, intent(out) :: K(2, 2)

    real*8 :: coef(2), dx(2), L, mb
    integer :: i, j
	L = abs(x(1) - x(2))
	dx = (/-1.0/L, 1.0/L/)
    if (tpCoupled == COMMON_1D) then
        coef = (sum(beta) + beta) * L / 6
        K(1, :) = dx(1) * coef
        K(2, :) = dx(2) * coef
    else if (tpCoupled == AXISYMMETRIC) then
        mB = sum(beta)*0.5D0
	    K(1, 1) = -(2*x(1) + x(2)) / 6 + 1.0D0 / 3 * L
	    K(1, 2) = -(x(1) + 2 * x(2)) / 6 + 1.0D0 / 6 * L
	    K(2, 1) = (2*x(1) + x(2)) / 6 + 1.0D0 / 6 * L
	    K(2, 2) = (x(1) + 2 * x(2)) / 6 + L / 3.0D0
	    K = mB * K
    end if
end subroutine


subroutine assembleHeatToElasticL2CTE_Func(this, beta, time, tpCoupled, A)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, external :: beta
	real*8, intent(in) :: time
    integer, intent(in) :: tpCoupled
    type(SMatrix) :: A

    integer :: i, j, idx(2), idx2(2), label
    real*8 :: x(2), bbeta(2), HeatToElastic(2, 2)
    type(Mesh1DL2), pointer :: mesh1D
    mesh1D => this%mesh1D
    do i = 1, this%nElement
        idx = mesh1D%elements(i)%idx
        label = mesh1D%elements(i)%num
        x = mesh1D%nodes(idx)%x(1)
        bbeta(1) = beta(x(1), label, time)
        bbeta(2) = beta(x(2), label, time)
        call elementCoupledMatrixL2CTE(x, bbeta, tpCoupled, HeatToElastic)
        call assembleMatrix(A, HeatToElastic, idx)
    end do
end subroutine

subroutine assembleHeatToElasticL2CTE_const(this, beta, tpCoupled, A)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in) :: beta
    integer, intent(in) :: tpCoupled
    type(SMatrix) :: A

    integer :: i, j, idx(2), idx2(2)
    real*8 :: x(2), bbeta(2), HeatToElastic(2, 2)
    type(Mesh1DL2), pointer :: mesh1D
    mesh1D => this%mesh1D
    bbeta = beta
    do i = 1, this%nElement
        idx = mesh1D%elements(i)%idx
        x = mesh1D%nodes(idx)%x(1)
        call elementCoupledMatrixL2CTE(x, bbeta, tpCoupled, HeatToElastic)
        call assembleMatrix(A, HeatToElastic, idx)
    end do
end subroutine

subroutine assembleHeatToElasticL2CTE_Vector(this, vecbeta, tpCoupled, A)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in) :: vecbeta(2)
    integer, intent(in) :: tpCoupled
    type(SMatrix) :: A

    integer :: i, j, idx(2), idx2(2)
    real*8 :: x(2), bbeta0(2), bbeta1(2), HeatToElastic(2, 2)
    type(Mesh1DL2), pointer :: mesh1D
    mesh1D => this%mesh1D
    bbeta0 = vecbeta(1)
	bbeta1 = vecbeta(2)
    do i = 1, this%nElement
        idx = mesh1D%elements(i)%idx
        x = mesh1D%nodes(idx)%x(1)
        call elementCoupledMatrixL2CTE_2(x, bbeta0, bbeta1, tpCoupled, HeatToElastic)
        call assembleMatrix(A, HeatToElastic, idx)
    end do
end subroutine

subroutine elementCoupledMatrixL2CTE_2(x, beta0, beta1, tpCoupled, K)
    implicit none
    real*8, intent(in) :: x(2), beta0(2), beta1(2)
    real*8, intent(out) :: K(2, 2)
    integer, intent(in) :: tpCoupled

    real*8 :: coef(2), mB1, mB2, dx(2), L
    integer :: i, j
	L = abs(x(1) - x(2))
	dx = (/-1.0/L, 1.0/L/)
    K = 0.0D0
    if (tpCoupled == AXISYMMETRIC) then
        mB1 = sum(beta0) * 0.5D0
	    mB2 = sum(beta1) * 0.5D0
	    K(1, 1) = -(2*x(1) + x(2)) / 6 * mB1 + 1.0D0 / 3 * L * mB2
	    K(1, 2) = -(x(1) + 2 * x(2)) / 6 * mB1 + 1.0D0 / 6 * L * mB2
	    K(2, 1) = (2*x(1) + x(2)) / 6 * mB1 + 1.0D0 / 6 * L * mB2
	    K(2, 2) = (x(1) + 2 * x(2)) / 6 * mB1 + L / 3.0D0 * mB2
    end if

end subroutine

subroutine assembleHeatSourceL2CTE_const(this, source, tp, RHS)
	implicit none
	class(FES1DL21CTE), intent(in) :: this
	real*8, intent(in) :: source
    integer, intent(in) :: tp
	real*8, intent(inout) :: RHS(this%nDof)

    integer :: i, j, index(2)
    real*8 :: x(2), coef(2), K(2)
    type(Mesh1DL2), pointer :: mesh1D
    mesh1D => this%mesh1D
    coef = source
    do i = 1, this%nElement
        index = mesh1D%elements(i)%idx
        x = mesh1D%nodes(index)%x(1)
		index = index + mesh1D%nv
        call elementHeatSource1DL21(x, coef, tp, K)
        call assembleVector(RHS, k, index)
    end do 
end subroutine

subroutine assembleDC1_L2CTE(this, rho, c, k, TRef, theta, E, Nu, beta, dt, &
                             tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in):: rho, c, k, E, Nu, beta, TRef, dt, theta
    type(SMatrix), intent(inout) :: LHS, RHS
    integer, intent(in) :: tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled       

    integer :: i, j, idx(2), idxHeat(2), idxElastic(2)
    real*8 :: x(2), coefM(2), coefK(2), coefE(2), coefNu(2), coefBeta(2), coefP(2)
    real*8 :: HeatStiff(2, 2), HeatMass(2, 2)
    real*8 :: ElasticStiff(2, 2), HeatToElastic(2, 2), ElasticToHeat(2, 2)
    type(mesh1DL2), pointer :: mesh0
    mesh0 => this%mesh1D
	coefP = rho
    coefM = rho * c / TRef
    coefK = k / TRef
    coefE = E
    coefNu = Nu
    coefBeta = beta
    do i = 1, mesh0%nt
        idx = mesh0%elements(i)%idx
        x = mesh0%nodes(idx)%x(1)
		idxElastic = idx
        idxHeat = idx + mesh0%nv
        call elementElasticityStiff1DL21(x, coefE, coefNu, tpElasticStiff, ElasticStiff)
        
        call assembleMatrix(LHS, ElasticStiff, idxElastic)
        
        call elementCoupledMatrixL2CTE(x, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)

        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)

        call elementHeatStiff1DL21(x, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass1DL21(x, coefM, tpHeatMass, HeatMass)
        call assembleMatrix(LHS, HeatMass + theta * dt * HeatStiff, idxHeat)
        call assembleMatrix(RHS, HeatMass - (1 - theta) * dt * HeatStiff, idxHeat)
    end do
end subroutine

subroutine assembleDC2_L2CTE_function(this, rho, c, k, E, Nu, beta, time, TRef, theta, dt, alpha, &
                                       tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, external :: rho, c, k, E, Nu, beta
	real*8, intent(in) :: TRef, time, dt, alpha, theta
    integer, intent(in) :: tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled
    type(SMatrix), intent(inout) :: LHS, RHS

    integer :: i, j, idx(2), label, idxHeat(2), idxElastic(2)
    real*8 :: x(2), coefM(2), coefK(2), coefE(2), coefNu(2), coefBeta(2), coefP(2)
    real*8 :: HeatStiff(2, 2), HeatMass(2, 2)
    real*8 :: ElasticStiff(2, 2), ElasticMass(2, 2)
    real*8 :: HeatToElastic(2, 2), ElasticToHeat(2, 2)

    type(mesh1DL2), pointer :: mesh0
    mesh0 => this%mesh1D

    do i = 1, mesh0%nt
        idx = mesh0%elements(i)%idx
        label = mesh0%elements(i)%num
        x = mesh0%nodes(idx)%x(1)
		do j = 1, 2
			coefP(j) = rho(x(j), label, time)
			coefM(j) = rho(x(j), label, time) * c(x(j), label, time) / TRef
			coefK(j) = k(x(j), label, time) / TRef
			coefE(j) = E(x(j), label, time)
			coefNu(j) = Nu(x(j), label, time)
			coefBeta(j) = beta(x(j), label, time)
		end do

		idxElastic = idx
        idxHeat = idx + mesh0%nv

		call elementELasticityStiff1DL21(x, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call elementHeatMass1DL21(x, coefP, tpElasticMass, ElasticMass)
        call elementCoupledMatrixL2CTE(x, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call elementHeatStiff1DL21(x, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass1DL21(x, coefM, tpHeatMass, HeatMass)

        call assembleMatrix(LHS, ElasticStiff + 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)     
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(LHS, HeatMass + theta * dt * HeatStiff, idxHeat)

        call assembleMatrix(RHS, 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, HeatMass - (1 - theta) * dt * HeatStiff, idxHeat)    
    end do
end subroutine


subroutine assembleDC2_L2CTE_Const(this, rho, c, k, TRef, theta, E, Nu, beta, dt, alpha, &
                                   tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in):: rho, c, k, TRef, theta, E, Nu, beta, dt, alpha
    integer, intent(in) :: tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled 
    type(SMatrix), intent(inout) :: LHS, RHS

    integer :: i, j, idx(2), idxHeat(2), idxElastic(2)
    real*8 :: x(2), coefM(2), coefK(2), coefE(2), coefNu(2), coefBeta(2), coefP(2)
    real*8 :: HeatStiff(2, 2), HeatMass(2, 2)
    real*8 :: ElasticStiff(2, 2), ElasticMass(2, 2)
    real*8 :: HeatToElastic(2, 2), ElasticToHeat(2, 2)

    type(mesh1DL2), pointer :: mesh0
    mesh0 => this%mesh1D

    coefE = E
    coefNu = Nu
    coefBeta = beta
    coefP = rho
	coefM = rho * c / TRef
	coefK = k / TRef
    do i = 1, mesh0%nt
        idx = mesh0%elements(i)%idx
        x = mesh0%nodes(idx)%x(1)
		idxElastic = idx
        idxHeat = idx + mesh0%nv

		call elementElasticityStiff1DL21(x, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call elementHeatMass1DL21(x, coefP, tpElasticMass, ElasticMass)
        call elementCoupledMatrixL2CTE(x, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call elementHeatStiff1DL21(x, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass1DL21(x, coefM, tpHeatMass, HeatMass)

        call assembleMatrix(LHS, ElasticStiff + 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)     
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(LHS, HeatMass + theta * dt * HeatStiff, idxHeat)

        call assembleMatrix(RHS, 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, HeatMass - (1 - theta) * dt * HeatStiff, idxHeat)    
    end do
end subroutine

subroutine assembleDC2_L2CTE_VM(this, rho, rhoC, k, D, beta0, beta1, TRef, dt, alpha, &
                                 tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in) :: rho, rhoC, k, D(2, 2), beta0, beta1
	real*8, intent(in) :: TRef, dt, alpha
    type(SMatrix), intent(inout) :: LHS, RHS
    integer, intent(in) :: tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled  

    integer :: i, j
    real*8 :: x(2), coefM(2), coefK(2), coefBeta0(2), coefBeta1(2), coefP(2)
    integer :: idx(2), idxHeat(2), idxElastic(2)
    real*8 :: HeatStiff(2, 2), HeatMass(2, 2)
    real*8 :: ElasticStiff(2, 2), ElasticMass(2, 2)
    real*8 :: HeatToElastic(2, 2), ElasticToHeat(2, 2)

    type(mesh1DL2), pointer :: mesh0
    mesh0 => this%mesh1D

    coefBeta0 = beta0
	coefBeta1 = beta1
    coefP = rho
	coefM = rhoC / TRef
	coefK = k / TRef
    do i = 1, mesh0%nt
        idx = mesh0%elements(i)%idx
        x = mesh0%nodes(idx)%x(1)
		idxElastic = idx
        idxHeat = idx + mesh0%nv

        call elementHeatStiff1DL21_M(x, D, tpElasticStiff, ElasticStiff)
        call elementHeatMass1DL21(x, coefP, tpElasticMass, ElasticMass)		
        call elementCoupledMatrixL2CTE_2(x, coefBeta0, coefBeta1, tpCoupled, HeatToElastic)

        ElasticToHeat = transpose(HeatToElastic)
        call elementHeatStiff1DL21(x, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass1DL21(x, coefM, tpHeatMass, HeatMass)

        call assembleMatrix(LHS, ElasticStiff + 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)     
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(LHS, HeatMass + 0.5 * dt * HeatStiff, idxHeat)

        call assembleMatrix(RHS, 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, HeatMass - 0.5 * dt * HeatStiff, idxHeat)    
    end do
end subroutine
                                 

subroutine extractUTL2CTE(this, x, U, T)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    real*8, intent(in) :: x(this%nDof)
    real*8, intent(out) :: U(this%nDof/2)
    real*8, intent(out) :: T(this%nDof/2)

    integer :: nv
    nv = this%nDof/2
    U = x(1:nv)
    T = x(nv+1:this%nDof)
end subroutine


subroutine applyBCZeroDispL2CTE(this, A, RHS, label)
    implicit none
    class(FES1DL21CTE), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(A%M)
    integer, intent(in) :: label(:)
    
    type(LinkedSet) :: dirichletBC
    integer, allocatable :: DOF(:)
    type(Mesh1DL2), pointer :: mesh1D
    mesh1D => this%mesh1D
    call mesh1D.extractBoundaryNodes(dirichletBC, label)
    allocate(DOF(dirichletBC%size))
    call getDofFromSet(dirichletBC, 1, DOF)
    call this.applyBCZeroDOF_MR(A, RHS, DOF)
    deallocate(DOF)
    call dirichletBC.empty()
end subroutine


subroutine applyBCHeatL2CTE_function(this, A, RHS, label, u0, time, tp)
	implicit none
	class(FES1DL21CTE), intent(in) :: this
	type(SMatrix), intent(inout) :: A
	real*8, intent(inout) :: RHS(this%nDof)
	integer, intent(in) :: label(:)
	real*8, external :: u0
    real*8, intent(in) :: time
	integer, intent(in) :: tp

	real*8, allocatable :: uu0(:)
	integer :: i
	real*8 :: x, y
	type(mesh1DL2), pointer :: mesh0
    type(LinkedSet) :: dirichletBC
	type(ele), pointer:: current_p
    mesh0 => this%mesh1D

	call mesh0.extractBoundaryNodes(dirichletBC, label)

	allocate(uu0(dirichletBC%size))

	i = 0
    current_p => dirichletBC%head
    do while(associated(current_p))
		i = i + 1
		x = mesh0%nodes(current_p%n)%x(1)
		uu0(i) = u0(x, time)
        current_p%n = current_p%n + mesh0%nv
        current_p => current_p%next
    end do
    call this.applyBC_SetDOF(A, RHS, dirichletBC, uu0, tp)
	deallocate(uu0)
	call dirichletBC.empty()
end subroutine

end module
